##########################################################################################

library(data.table)
library(optparse)
library(pheatmap)
library(cowplot)
library(GSVA)

##########################################################################################

option_list <- list(
    make_option(c("--geneset_type"), type = "character"),
    make_option(c("--geneset_file"), type = "character"),
    make_option(c("--rna_file"), type = "character"),
    make_option(c("--motif_file"), type = "character"),
    make_option(c("--cor_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

    geneset_type <- "known_motif"
    rna_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/GeneExpression.MeanByCellType.tsv"
    motif_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/Motif.MeanByCellType.tsv"
    cor_file <- "~/20231121_singleMuti/results/celltype_plot/mfuzz/cor.motif_atac-rna.tsv"
    out_path <- "~/20231121_singleMuti/results/celltype_plot/pheatmap"
    geneset_file <- "~/20231121_singleMuti/config/Human_reported_TF2.csv"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

rna_file <- opt$rna_file
motif_file <- opt$motif_file
cor_file <- opt$cor_file
out_path <- opt$out_path
geneset_file <- opt$geneset_file
geneset_type <- opt$geneset_type

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
sco_exp <- data.frame(fread(rna_file))
sco_motif <- data.frame(fread(motif_file))
dat_cor <- data.frame(fread(cor_file))
dat_geneset <- data.frame(fread(geneset_file , header = F))

##########################################################################################

cell_order <- c("SSC" , "Differenting.Differented.SPG" , "Leptotene" ,
    "Zygotene" , "Patchytene" , "Diplotene" , 
    "Early.stage.of.spermatids" , "Round.ElongateS.tids" , "Sperm"
    )

##########################################################################################

rownames(sco_exp) <- sco_exp$gene
sco_exp <- sco_exp[,-1]

rownames(sco_motif) <- sco_motif$gene
sco_motif <- sco_motif[,-1]

##########################################################################################
## 提取感兴趣的基因集合

tf_name <- unique(dat_geneset$V1)
tf_name <- tf_name[tf_name %in% dat_cor$GeneExpressionMatrix_name]

sco_exp <- sco_exp[tf_name,]
sco_motif <- sco_motif[tf_name,]

##########################################################################################

sco_exp <- sco_exp[,cell_order]
sco_motif <- sco_motif[,cell_order]

##########################################################################################
# 定义函数将矩阵的每一行标准化到-1到1的范围
row_normalize <- function(matrix_data) {
  normalized_matrix <- apply(matrix_data, 1, function(x) {
    min_val <- min(x)
    max_val <- max(x)
    normalized_row <- ((x - min_val) / (max_val - min_val)) * 2 - 1
    return(normalized_row)
  })
  return(normalized_matrix)
}

normalized_data_sco_exp <- t(row_normalize(sco_exp))
normalized_data_sco_motif <- t(row_normalize(sco_motif))
normalized_data_sco_exp_motif <- normalized_data_sco_exp - normalized_data_sco_motif

## 按表达的高低排
gene_order <- rownames(normalized_data_sco_exp)[order( normalized_data_sco_exp[,1] , normalized_data_sco_exp[,2] , normalized_data_sco_exp[,3] ,
    normalized_data_sco_exp[,4] , normalized_data_sco_exp[,5] , normalized_data_sco_exp[,6] ,
    normalized_data_sco_exp[,7] , normalized_data_sco_exp[,8] , normalized_data_sco_exp[,9] , decreasing = T
    )]

## 按表达-motif的高低排
#gene_order <- rownames(normalized_data_sco_exp_motif)[order( normalized_data_sco_exp_motif[,1] , normalized_data_sco_exp_motif[,2] , normalized_data_sco_exp_motif[,3] ,
#    normalized_data_sco_exp_motif[,4] , normalized_data_sco_exp_motif[,5] , normalized_data_sco_exp_motif[,6] ,
#    normalized_data_sco_exp_motif[,7] , normalized_data_sco_exp_motif[,8] , normalized_data_sco_exp_motif[,9] , decreasing = T
#    )]


## 按相关系数的高低排
gene_order <- dat_cor[order(dat_cor$cor , decreasing = T),"GeneExpressionMatrix_name"]
gene_order <- gene_order[ gene_order %in% tf_name ]

for( clus in seq(2:round(length(tf_name)/2)) ){

    ## 按照exp排序
    p0 <- pheatmap(normalized_data_sco_exp[gene_order,],
        #scale = "row", #annotation_row = gene_order ,
        show_rownames = T ,show_colnames = T, 
        #cutree_cols=7,cutree_rows = 7,
        cluster_rows = T , cluster_cols = F , clustering_method = "ward.D2",
        color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
        cutree_rows = clus ,
        cellheight = 10)
    col_cluster <- cutree(p0$tree_row,k=clus)
    gsva_path <- data.frame( gene = names(col_cluster) , group_name = col_cluster )
    gene_order <- gsva_path[order(gsva_path$group_name , decreasing = T),"gene"]

    order <- p0$tree_row$order
    gene_order <- row.names(normalized_data_sco_exp[gene_order,])[order]

    p <- pheatmap(normalized_data_sco_exp[gene_order,],
        #scale = "row", #annotation_row = gene_order ,
        show_rownames = T ,show_colnames = T, 
        #cutree_cols=7,cutree_rows = 7,
        cluster_rows = F , cluster_cols = F , clustering_method = "ward.D2",
        color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
        #cutree_rows = 9 ,
        cellheight = 10)

    p2 <- pheatmap(normalized_data_sco_motif[gene_order,],
        #scale = "row", #annotation_row = gene_order ,
        show_rownames = T ,show_colnames = T, 
        #cutree_cols=7,cutree_rows = 7,
        cluster_rows = F , cluster_cols = F , clustering_method = "ward.D2",
        color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
        #cutree_rows = 9 ,
        cellheight = 10)

    p3 <- pheatmap(normalized_data_sco_exp_motif[gene_order,],
        #scale = "row", #annotation_row = gene_order ,
        show_rownames = T ,show_colnames = T, 
        #cutree_cols=7,cutree_rows = 7,
        cluster_rows = F , cluster_cols = F , clustering_method = "ward.D2",
        color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
        #cutree_rows = 9 ,
        cellheight = 10)

    out_file <- paste0( out_path , "/pheatmap_tf." , geneset_type , "." , clus , ".pdf" )
    pdf(out_file , height = 20 , width = 10)
    #print(p + p2 + p3)
    print(p0)
    print(cowplot::plot_grid(p$gtable , p2$gtable , p3$gtable , ncol= 3, labels=c("Exp" , "Motif" , "Exp-Motif")))
    dev.off()
}